package queue

import (
	"fmt"
	"github.com/RichardKnop/machinery/v2"
	redisbackend "github.com/RichardKnop/machinery/v2/backends/redis"
	amqpbroker "github.com/RichardKnop/machinery/v2/brokers/amqp"
	"github.com/RichardKnop/machinery/v2/config"
	eagerlock "github.com/RichardKnop/machinery/v2/locks/eager"
	"github.com/RichardKnop/machinery/v2/log"
	"github.com/RichardKnop/machinery/v2/tasks"
	uuid "github.com/satori/go.uuid"
	"github.com/sirupsen/logrus"
	"github.com/spf13/viper"
)

var (
	AsyncTaskCenter *machinery.Server
)

func InitQueueWorkConfig() {
	tc, err := NewTaskCenter()
	if err != nil {
		panic(err)
	}
	AsyncTaskCenter = tc
}

// 获取amqp连接地址
func getAmqpUrl() (amqpURI string) {

	host := viper.GetString("amqp.host")
	port := viper.GetString("amqp.port")
	username := viper.GetString("amqp.user")
	password := viper.GetString("amqp.password")

	amqpURI = fmt.Sprintf("amqp://%s:%s@%s:%s/",
		username,
		password,
		host,
		port)

	return
}

// 获取redis连接地址 [password@]host[port][/db_num]
func getRedisUrl() (redisUrl string, db int) {
	host := viper.GetString("redis.host")
	port := viper.GetString("redis.port")
	password := viper.GetString("redis.password")
	db = viper.GetInt("redis.select")

	redisUrl = fmt.Sprintf("%s@%s:%s",
		password,
		host,
		port)

	return
}

func NewTaskCenter() (*machinery.Server, error) {

	// 获取rabbitmq连接地址
	amqpURI := getAmqpUrl()

	// 获取redis连接地址
	redisUrl, redisDB := getRedisUrl()

	transcodeName := viper.GetString("amqp.queueName")
	if transcodeName == "" {
		transcodeName = "transcode"
	}

	var queueName string
	queueAutoDeleted := viper.GetBool("amqp.queueAutoDeleted")
	if queueAutoDeleted {
		uuidStr := uuid.NewV4().String()
		queueName = transcodeName + "_queue_" + uuidStr
	} else {
		queueName = transcodeName + "_queue"
	}

	resultsExpireIn := viper.GetInt("amqp.resultsExpireIn")

	if !(resultsExpireIn > 0) {
		resultsExpireIn = 3600
	}
	cnf := &config.Config{
		Broker:          amqpURI,
		DefaultQueue:    queueName,
		ResultBackend:   redisUrl, // 使用 redis 保存任务的状态和结果集
		ResultsExpireIn: resultsExpireIn,
		NoUnixSignals:   true,
		AMQP: &config.AMQPConfig{
			Exchange:      transcodeName + "_exchange",
			ExchangeType:  "direct",
			BindingKey:    transcodeName + "_key", // 队列用这个键绑定到交换器
			PrefetchCount: 1,                      // 预取多少任务(如果有长时间运行的任务，设置为1)
			AutoDelete:    queueAutoDeleted,
		},
	}

	// 创建服务实体
	broker := amqpbroker.New(cnf)
	backend := redisbackend.NewGR(cnf, []string{redisUrl}, redisDB)
	lock := eagerlock.New()
	server := machinery.NewServer(cnf, broker, backend, lock)

	initAsyncTaskMap()
	return server, server.RegisterTasks(asyncTaskMap)
}

func NewAsyncTaskWorker(consumerTag string, concurrency int) *machinery.Worker {

	// The second argument is a consumer tag
	// Ideally, each worker should have a unique tag (worker1, worker2 etc)
	worker := AsyncTaskCenter.NewWorker(consumerTag, concurrency)

	// Here we inject some custom code for error handling,
	// start and end of task hooks, useful for metrics for example.
	errorhandler := func(err error) {
		logrus.WithFields(logrus.Fields{
			"route": "errorhandler error",
			"request_param": logrus.Fields{
				"Param": "",
			},
		}).Error(err.Error())
	}
	pretaskhandler := func(signature *tasks.Signature) {
		log.INFO.Println("I am a start of task handler for:", signature.Name)
	}
	posttaskhandler := func(signature *tasks.Signature) {
		log.INFO.Println("I am an end of task handler for:", signature.Name)
	}

	worker.SetPostTaskHandler(posttaskhandler)
	worker.SetErrorHandler(errorhandler)
	worker.SetPreTaskHandler(pretaskhandler)
	return worker
}
